import torch
from torch import nn, optim
from torchvision import datasets, transforms

filepath = '/home/h/筑工/data/'


# 训练集的数据预处理
transform_train = transforms.Compose([
    # 数据增强，随机裁剪224*224大小
    transforms.RandomResizedCrop(224),
    # 数据增强，随机水平翻转
    transforms.RandomHorizontalFlip(),
    # 数据变成tensor类型，像素值归一化，调整维度[h,w,c]==>[c,h,w]
    transforms.ToTensor(),
    # 对每个通道的像素进行标准化，给出每个通道的均值和方差
    transforms.Normalize(mean=(0.5,0.5,0.5), std=(0.5,0.5,0.5))])
 
# 验证集的数据预处理
transform_val = transforms.Compose([
    # 将输入图像大小调整为224*224
    transforms.Resize((224,224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=(0.5,0.5,0.5), std=(0.5,0.5,0.5))])

# 读取训练集并预处理
train_dataset = datasets.ImageFolder(root=filepath + 'train',  # 训练集图片所在的文件夹
                                     transform = transform_train)  # 训练集的预处理方法

 
# 读取验证集并预处理
val_dataset = datasets.ImageFolder(root=filepath + 'val',  # 验证集图片所在的文件夹
                                     transform = transform_val)  # 验证集的预处理方法

# 查看训练集和验证集的图片数量
train_num = len(train_dataset)
val_num = len(val_dataset)
print('train_num:', train_num, 'val_num:', val_num)   # 453, 112
 
# 查看图像类别及其对应的索引
class_dict = train_dataset.class_to_idx
print(class_dict)  # {'Bananaquit': 0, 'Black Skimmer': 1, 'Black Throated Bushtiti': 2, 'Cockatoo': 3}
# 将类别名称保存在列表中
class_names = list(class_dict.keys())